'''
Created on Oct 9, 2009

@author: mkiyer
'''

from veggie.app.alignment.bowtie import Bowtie
from veggie.sample.sampledb2 import get_sampledb
import operator
import sys
import logging
import os
import glob
import re
import collections

import numpy as np
import matplotlib.pyplot as plt
    
organism_re = re.compile(r"""organism=(.+?);""")

def read_viruses_gff(fhd):
    for thisline in fhd:
        thisline = thisline.strip()
        if not thisline:
            continue
        if thisline.startswith('#'):
            continue

        thisfields = thisline.split('\t')
        record_type = thisfields[2]
        if record_type == 'source':
            accession = thisfields[0]
            info = thisfields[8]
            organism = organism_re.match(info).group(1)            
            yield accession, organism, info

def parse_bowtie_virus_results(fhd):
    for thisline in fhd:
        thisline = thisline.strip()
        if not thisline:
            continue
        thisfields = thisline.split('\t')
        ref = thisfields[2]
        ref_fields = ref.split('|')
        accession = ref_fields[1]        
#        pos_fields = ref_fields[2].split('-')
#        start = int(pos_fields[0][1:])
#        end = int(pos_fields[1])
        yield accession

def count_virus_reads(fhd):
    virus_reads = collections.defaultdict(lambda: 0)
    for acc in parse_bowtie_virus_results(fhd):
        virus_reads[acc] += 1
    sorted_virus_reads = sorted(virus_reads.iteritems(), 
                                key=operator.itemgetter(1), 
                                reverse=False)
    return sorted_virus_reads


def plot_barplot(virus_reads, accession_lookup,
                 output_file=None, title=""):
    num_bars = len(virus_reads)
    read_counts = np.array([x[1] for x in virus_reads])
    accessions = [x[0] for x in virus_reads]

    fig = plt.figure(figsize=(10,8))
    ax1 = fig.add_subplot(111)
    plt.subplots_adjust(left=0.15, right=0.95)
    fig.canvas.set_window_title('Virus Read Counts')
    pos = np.arange(num_bars)+0.5    #Center bars on the Y-axis ticks
    rects = ax1.barh(pos, read_counts, align='center', height=0.8, color='y')
    plt.yticks(pos, accessions)

    # attach some text labels
    for i, rect in enumerate(rects):
        width = rect.get_width()
        plt.text(0.1, rect.get_y() + 0.05,
                 '%s=%d'% (accession_lookup[accessions[i]], int(width)),
                 color='b', size='small',
                 ha='left', va='bottom')

    plt.title('%s read counts' % title)
    plt.grid(True)
    if output_file is None:
        plt.show()
    else:
        plt.savefig(output_file)
    return

if __name__ == '__main__':
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)
    
    bowtie_res_file = sys.argv[1]
    name = os.path.splitext(os.path.basename(bowtie_res_file))[0]
    png_file = name + '.png'

    logging.debug("building accession lookup table")
    accession_lookup = {}    
    for res in read_viruses_gff(open('viruses.gff')):
        accession, organism, info = res
        accession_lookup[accession] = organism

    virus_reads = count_virus_reads(open(bowtie_res_file))
    plot_barplot(virus_reads, accession_lookup, output_file=png_file, title=name)
    for k,v in virus_reads:
        output_string = '\t'.join(map(str,[k, v, accession_lookup[k]]))
        print output_string
        
    
    
    
        
